#!/usr/bin/env python
"""
FRP Client implementation for TCP port forwarding.
This module implements a client that connects to an FRP server and forwards traffic
to a target host and port.
"""

import socket
import time
import threading
import struct
import selectors
import logging
from typing import List, Optional
import lib.ConnTool as ConnTool
import os
from datetime import datetime

# pyinstaller --onefile frpc.py

log_dir = "logs"
os.makedirs(log_dir, exist_ok=True)
log_file = os.path.join(log_dir, f"frpc_{datetime.now().strftime('%Y%m%d')}.log")

# 创建日志格式化器
formatter = logging.Formatter("%(asctime)s [%(levelname)s] %(message)s")

# 创建文件处理器
file_handler = logging.FileHandler(log_file, encoding="utf-8")
file_handler.setFormatter(formatter)

# 创建控制台处理器
console_handler = logging.StreamHandler()
console_handler.setFormatter(formatter)

# 获取根日志记录器
logger = logging.getLogger()
logger.addHandler(file_handler)
logger.addHandler(console_handler)  # 添加控制台处理器
logger.setLevel(logging.INFO)

# 创建全局选择器
selector = selectors.DefaultSelector()


class Frpc:
    """
    FRP Client class that handles connection to FRP server and port forwarding.

    Attributes:
        target_host (str): The target host to forward traffic to
        target_port (int): The target port to forward traffic to
        server_host (str): The FRP server host to connect to
        server_port (int): The FRP server port to connect to
        server_fd (socket.socket): Socket connection to the FRP server
        work_conn_pool (List[socket.socket]): Pool of work connections
        is_connected (bool): Whether the client is connected to the server
        reconnect_delay (int): Reconnection delay time in seconds
        max_reconnect_delay (int): Maximum reconnection delay time in seconds
        _running (bool): Running state flag
    """

    def __init__(
        self, server_host: str, server_port: int, target_host: str, target_port: int
    ):
        """
        Initialize the FRP client.

        Args:
            server_host: FRP server host address
            server_port: FRP server port number
            target_host: Target host address to forward to
            target_port: Target port number to forward to
        """
        self.target_host = target_host
        self.target_port = target_port
        self.server_host = server_host
        self.server_port = server_port
        self.server_fd: Optional[socket.socket] = None
        self.work_conn_pool: List[socket.socket] = []
        self.is_connected = False
        self.reconnect_delay = 5  # 初始重连延迟为5秒
        self.max_reconnect_delay = 30  # 最大重连延迟为30秒
        self._running = True  # 添加运行状态标志

        # 启动连接
        self.connect_to_server()

        # 启动维护线程
        self.maintain_thread = threading.Thread(
            target=self.maintain_conn_pool, daemon=True
        )
        self.heartbeat_thread = threading.Thread(target=self.heartbeat, daemon=True)
        self.maintain_thread.start()
        self.heartbeat_thread.start()

    def connect_to_server(self) -> None:
        """建立与服务器的连接"""
        try:
            if self.server_fd is not None:
                try:
                    # 先注销旧的 socket
                    selector.unregister(self.server_fd)
                    self.server_fd.close()
                except:
                    pass
                self.server_fd = None

            self.server_fd = socket.create_connection(
                (self.server_host, self.server_port)
            )
            self.server_fd.sendall(struct.pack("i", 1))  # 发送初始心跳包
            self.server_fd.setblocking(False)
            selector.register(
                self.server_fd, selectors.EVENT_READ, self.handle_controller_data
            )
            self.is_connected = True
            self.reconnect_delay = 5  # 重置重连延迟
            logger.info("已连接到服务器")
        except Exception as e:
            self.is_connected = False
            logger.error(f"连接服务器失败: {e}")
            raise

    def stop(self):
        """停止FRP客户端的所有活动"""
        self._running = False  # 停止所有线程
        self.is_connected = False

        # 发送断开连接命令给服务器
        if self.server_fd:
            try:
                # 发送断开连接命令(3)给服务器
                self.server_fd.send(struct.pack("i", 3))
                logger.info("已发送断开连接命令")
                # 等待一小段时间确保命令发送成功
                time.sleep(0.2)
            except Exception as e:
                logger.error(f"发送断开连接命令失败: {e}")

        # 关闭所有工作连接
        for conn in self.work_conn_pool:
            try:
                conn.close()
            except:
                pass
        self.work_conn_pool.clear()

        # 关闭服务器连接
        if self.server_fd:
            try:
                self.server_fd.close()
            except:
                pass
            self.server_fd = None

    def heartbeat(self) -> None:
        """向服务器发送定期心跳包"""
        while self._running:  # 使用运行状态标志
            try:
                if not self.is_connected:
                    if not self._running:  # 检查是否应该停止
                        break
                    logger.info("尝试重新连接服务器...")
                    try:
                        self.connect_to_server()
                    except:
                        if not self._running:  # 检查是否应该停止
                            break
                        # 如果重连失败，增加延迟时间
                        self.reconnect_delay = min(
                            self.reconnect_delay * 2, self.max_reconnect_delay
                        )
                        logger.info(f"重连失败，{self.reconnect_delay}秒后重试")
                        time.sleep(self.reconnect_delay)
                        continue

                if self.server_fd is not None:
                    try:
                        self.server_fd.send(struct.pack("i", 1))
                        logger.debug("发送心跳包")
                    except Exception as e:
                        logger.error(f"心跳包发送错误: {e}")
                        self.is_connected = False
                        continue

                time.sleep(9)
            except Exception as e:
                logger.error(f"心跳线程错误: {e}")
                self.is_connected = False
                time.sleep(1)

    def maintain_conn_pool(self) -> None:
        """维护连接池"""
        logger.info("启动TCP连接池")
        pool_size = 0
        while self._running:  # 使用运行状态标志
            try:
                if not self.is_connected:
                    time.sleep(1)
                    continue

                if len(self.work_conn_pool) < pool_size:
                    work_conn = socket.create_connection(
                        (self.server_host, self.server_port)
                    )
                    target_conn = socket.create_connection(
                        (self.target_host, self.target_port)
                    )
                    ConnTool.join(target_conn, work_conn)
                    self.work_conn_pool.append(work_conn)
            except Exception as e:
                logger.error(f"连接池维护错误: {e}")
                time.sleep(1)

    def handle_controller_data(self, server_fd: socket.socket, mask: int) -> None:
        """
        Handle control data received from the server.

        Args:
            server_fd: The server socket file descriptor
            mask: The event mask
        """
        try:
            data = server_fd.recv(4)  # 接收4字节的命令数据
            if data:
                cmd = struct.unpack("i", data)[0]
                logger.debug(f"收到命令: {cmd}")

                if cmd == 2:  # 建立工作TCP连接的命令
                    logger.info("收到frps控制指令")
                    if self.work_conn_pool:
                        work_conn = self.work_conn_pool.pop()
                    else:
                        target_conn = socket.create_connection(
                            (self.target_host, self.target_port)
                        )
                        work_conn = socket.create_connection(
                            (self.server_host, self.server_port)
                        )
                        ConnTool.join(target_conn, work_conn)

                    work_conn.sendall(struct.pack("i", 2))  # 发送确认包
                    logger.info("建立工作TCP连接")

        except IOError:
            # 非阻塞模式下，如果没有数据可读会抛出异常
            pass
        except Exception as e:
            logger.error(f"处理控制数据错误: {e}")
            self.is_connected = False

    def run(self) -> None:
        """Main event loop for handling server events."""
        try:
            while self._running:  # 使用运行状态标志
                if not self.is_connected:
                    time.sleep(1)
                    continue

                events = selector.select()
                for key, mask in events:
                    callback = key.data
                    callback(key.fileobj, mask)
        except KeyboardInterrupt:
            logger.info("程序被用户中断")
        except Exception as e:
            logger.error(f"运行错误: {e}")
        finally:
            logger.info("frpc已停止")


def main():
    """Main entry point for the FRP client."""
    logger.info("正在启动frpc...")

    # ============= 配置参数 =============
    # 远程服务器配置
    remote_host = "xxxxx"  # FRP服务器地址
    remote_port = xxxxx  # FRP服务器端口

    # 目标服务器配置
    target_host = "172.16.3.70"  # 目标服务器地址
    target_port = 7001  # 目标服务器端口
    # ===================================

    try:
        logger.info(f"FRP服务器地址: {remote_host}")
        logger.info(f"FRP服务器端口: {remote_port}")
        logger.info(f"目标服务器地址: {target_host}")
        logger.info(f"目标服务器端口: {target_port}")
        client = Frpc(remote_host, remote_port, target_host, target_port)
        client.run()
    except Exception as e:
        logger.error(f"启动失败: {e}")


if __name__ == "__main__":
    main()
